#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@file    statics.py
@brief
@details
@author  Shivelino
@date    2023-12-23 19:16
@version 0.0.1

@par Copyright(c):
@par todo:
@par history:
"""
import argparse
import os.path as op
import matplotlib.pyplot as plt
import pickle as pkl
import numpy as np

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default="mobilenet", help='model')
    parser.add_argument('--data_dir', type=str, default="../out", help='data directory')
    parser.add_argument('--loss', type=bool, default=True, help='loss statics')
    parser.add_argument('--acc', type=bool, default=True, help='accuracy statics')
    parser.add_argument('--save_img', type=bool, default=True, help='save statics image')
    opt = parser.parse_args()

    if opt.loss:
        losses = pkl.load(open(op.join(opt.data_dir, f"losses_{opt.model}.pkl"), "rb"))
        x_axis = np.array(range(1, len(losses) + 1))
        plt.title(f"{opt.model} Loss")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.xticks(x_axis)
        plt.plot(x_axis, losses, label="loss")
        if opt.save_img:
            plt.savefig(op.join(opt.data_dir, f"loss_{opt.model}.jpg"))
        plt.show()

    if opt.acc:
        acces = pkl.load(open(op.join(opt.data_dir, f"acces_{opt.model}.pkl"), "rb"))
        x_axis = np.array(range(1, len(acces) + 1))
        plt.title(f"{opt.model} Accuracy")
        plt.xlabel("Epoch")
        plt.ylabel("Accuracy")
        plt.xticks(x_axis)
        plt.plot(x_axis, acces, label="accuracy")
        if opt.save_img:
            plt.savefig(op.join(opt.data_dir, f"accuracy_{opt.model}.jpg"))
        plt.show()
